package portforward

import (
	"context"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"time"
)

// 在 PortForwarder 结构体中添加跟踪监听器
type PortForwarder struct {
	mu        sync.Mutex
	proxies   map[string]context.CancelFunc
	listeners map[string]net.Listener   // 跟踪 TCP 监听器
	udpConns  map[string]net.PacketConn // 跟踪 UDP 连接
	waitGroup sync.WaitGroup
	config    *ForwardingConfig
	logger    Logger
}

// 修改 NewPortForwarder
func NewPortForwarder(config *ForwardingConfig, logger Logger) *PortForwarder {
	if config == nil {
		config = DefaultConfig()
	}
	if logger == nil {
		logger = &defaultLogger{}
	}
	return &PortForwarder{
		proxies:   make(map[string]context.CancelFunc),
		listeners: make(map[string]net.Listener),
		udpConns:  make(map[string]net.PacketConn),
		config:    config,
		logger:    logger,
	}
}

// ForwardingConfig 存储转发配置
type ForwardingConfig struct {
	BufferSize    int           // 缓冲区大小
	ReadTimeout   time.Duration // 读取超时时间
	WriteTimeout  time.Duration // 写入超时时间
	MaxRetries    int           // 最大重试次数
	RetryInterval time.Duration // 重试间隔
}

// DefaultConfig 返回默认配置
// func DefaultConfig() *ForwardingConfig {
// 	return &ForwardingConfig{
// 		BufferSize:    65535,
// 		ReadTimeout:   30 * time.Second,
// 		WriteTimeout:  30 * time.Second,
// 		MaxRetries:    3,
// 		RetryInterval: 1 * time.Second,
// 	}
// }

func DefaultConfig() *ForwardingConfig {
	return &ForwardingConfig{
		BufferSize:    32 * 1024,        // 32KB buffer
		ReadTimeout:   2 * time.Minute,  // 增加超时时间到2分钟
		WriteTimeout:  30 * time.Second, // 写入超时保持30秒
		MaxRetries:    3,
		RetryInterval: time.Second,
	}
}

// Logger 定义日志接口
type Logger interface {
	Info(format string, args ...interface{})
	Error(format string, args ...interface{})
}

// defaultLogger 实现默认的日志记录
type defaultLogger struct{}

func (l *defaultLogger) Info(format string, args ...interface{}) {
	fmt.Printf("[INFO] "+format+"\n", args...)
}

func (l *defaultLogger) Error(format string, args ...interface{}) {
	fmt.Printf("[ERROR] "+format+"\n", args...)
}

// StartTCPProxy 启动 TCP 端口转发
func (pf *PortForwarder) StartTCPProxy(listenAddr, targetAddr string) error {
	pf.mu.Lock()
	defer pf.mu.Unlock()

	if _, exists := pf.proxies[listenAddr]; exists {
		return fmt.Errorf("TCP forwarding already exists: %s", listenAddr)
	}

	ctx, cancel := context.WithCancel(context.Background())
	listener, err := net.Listen("tcp", listenAddr)
	if err != nil {
		cancel()
		return fmt.Errorf("failed to listen (TCP): %v", err)
	}

	pf.proxies[listenAddr] = cancel
	pf.listeners[listenAddr] = listener // 保存监听器引用
	pf.waitGroup.Add(1)

	go func() {
		defer pf.waitGroup.Done()
		defer listener.Close()

		pf.logger.Info("Listening (TCP) %s, forwarding to %s", listenAddr, targetAddr)
		for {
			client, err := listener.Accept()
			if err != nil {
				select {
				case <-ctx.Done():
					return
				default:
					// 检查是否是因为监听器关闭导致的错误
					if ne, ok := err.(net.Error); ok && ne.Temporary() {
						pf.logger.Error("Temporary accept error (TCP): %v", err)
						time.Sleep(time.Second)
						continue
					}
					pf.logger.Error("Accept error (TCP): %v", err)
					return
				}
			}
			go pf.handleTCPConnection(ctx, client, targetAddr)
		}
	}()
	return nil
}

// StartUDPProxy 启动 UDP 端口转发
// 修改 StartUDPProxy
func (pf *PortForwarder) StartUDPProxy(listenAddr, targetAddr string) error {
	pf.mu.Lock()
	defer pf.mu.Unlock()

	if _, exists := pf.proxies[listenAddr]; exists {
		return fmt.Errorf("UDP forwarding already exists: %s", listenAddr)
	}

	ctx, cancel := context.WithCancel(context.Background())
	listenConn, err := net.ListenPacket("udp", listenAddr)
	if err != nil {
		cancel()
		return fmt.Errorf("failed to listen (UDP): %v", err)
	}

	pf.proxies[listenAddr] = cancel
	pf.udpConns[listenAddr] = listenConn // 保存 UDP 连接引用
	pf.waitGroup.Add(1)

	go func() {
		defer pf.waitGroup.Done()
		defer listenConn.Close()

		pf.logger.Info("Listening (UDP) %s, forwarding to %s", listenAddr, targetAddr)

		// 使用 map 来跟踪客户端连接
		type clientInfo struct {
			targetConn net.Conn
			lastActive time.Time
		}
		clients := make(map[string]*clientInfo)
		var clientsMu sync.RWMutex

		// 清理过期连接
		go func() {
			ticker := time.NewTicker(time.Minute)
			defer ticker.Stop()
			for {
				select {
				case <-ctx.Done():
					return
				case <-ticker.C:
					now := time.Now()
					clientsMu.Lock()
					for addr, client := range clients {
						if now.Sub(client.lastActive) > pf.config.ReadTimeout {
							client.targetConn.Close()
							delete(clients, addr)
						}
					}
					clientsMu.Unlock()
				}
			}
		}()

		buffer := make([]byte, pf.config.BufferSize)
		for {
			n, clientAddr, err := listenConn.ReadFrom(buffer)
			if err != nil {
				select {
				case <-ctx.Done():
					return
				default:
					pf.logger.Error("Failed to read UDP data: %v", err)
					continue
				}
			}

			clientKey := clientAddr.String()
			clientsMu.RLock()
			client, exists := clients[clientKey]
			clientsMu.RUnlock()

			if !exists {
				targetConn, err := net.DialTimeout("udp", targetAddr, pf.config.WriteTimeout)
				if err != nil {
					pf.logger.Error("Failed to connect to target (UDP): %v", err)
					continue
				}

				client = &clientInfo{
					targetConn: targetConn,
					lastActive: time.Now(),
				}

				clientsMu.Lock()
				clients[clientKey] = client
				clientsMu.Unlock()

				// 处理目标服务器的响应
				go func() {
					responseBuffer := make([]byte, pf.config.BufferSize)
					for {
						client.targetConn.SetReadDeadline(time.Now().Add(pf.config.ReadTimeout))
						n, err := client.targetConn.Read(responseBuffer)
						if err != nil {
							if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
								continue
							}
							return
						}

						clientsMu.Lock()
						client.lastActive = time.Now()
						clientsMu.Unlock()

						_, err = listenConn.WriteTo(responseBuffer[:n], clientAddr)
						if err != nil {
							pf.logger.Error("Failed to write response to client (UDP): %v", err)
							return
						}
					}
				}()
			}

			clientsMu.Lock()
			client.lastActive = time.Now()
			clientsMu.Unlock()

			client.targetConn.SetWriteDeadline(time.Now().Add(pf.config.WriteTimeout))
			_, err = client.targetConn.Write(buffer[:n])
			if err != nil {
				pf.logger.Error("Failed to write to target (UDP): %v", err)
				continue
			}
		}
	}()
	return nil
}

// handleTCPConnection 处理 TCP 连接
func (pf *PortForwarder) handleTCPConnection(ctx context.Context, client net.Conn, target string) {
	defer client.Close()

	var targetConn net.Conn
	var err error

	// 连接目标服务器
	for i := 0; i < pf.config.MaxRetries; i++ {
		targetConn, err = net.DialTimeout("tcp", target, pf.config.WriteTimeout)
		if err == nil {
			break
		}
		if i < pf.config.MaxRetries-1 {
			pf.logger.Info("Retrying connection to target (%d/%d)", i+1, pf.config.MaxRetries)
			time.Sleep(pf.config.RetryInterval)
		}
	}
	if err != nil {
		pf.logger.Error("Failed to connect to target (TCP): %v", err)
		return
	}
	defer targetConn.Close()

	// 创建一个子上下文，用于控制数据传输
	transferCtx, cancelTransfer := context.WithCancel(ctx)
	defer cancelTransfer()

	// 定期更新超时时间
	go func() {
		ticker := time.NewTicker(pf.config.ReadTimeout / 2)
		defer ticker.Stop()

		for {
			select {
			case <-transferCtx.Done():
				return
			case <-ticker.C:
				// 更新读写超时
				deadline := time.Now().Add(pf.config.ReadTimeout)
				client.SetDeadline(deadline)
				targetConn.SetDeadline(deadline)
			}
		}
	}()

	// 错误通道
	errChan := make(chan error, 2)

	// 客户端到目标的数据传输
	go func() {
		buf := make([]byte, pf.config.BufferSize)
		for {
			select {
			case <-transferCtx.Done():
				return
			default:
				n, err := client.Read(buf)
				if err != nil {
					if err != io.EOF && !isConnectionClosed(err) {
						errChan <- fmt.Errorf("client read error: %v", err)
					}
					return
				}

				if n > 0 {
					_, err = targetConn.Write(buf[:n])
					if err != nil {
						if !isConnectionClosed(err) {
							errChan <- fmt.Errorf("target write error: %v", err)
						}
						return
					}
				}
			}
		}
	}()

	// 目标到客户端的数据传输
	go func() {
		buf := make([]byte, pf.config.BufferSize)
		for {
			select {
			case <-transferCtx.Done():
				return
			default:
				n, err := targetConn.Read(buf)
				if err != nil {
					if err != io.EOF && !isConnectionClosed(err) {
						errChan <- fmt.Errorf("target read error: %v", err)
					}
					return
				}

				if n > 0 {
					_, err = client.Write(buf[:n])
					if err != nil {
						if !isConnectionClosed(err) {
							errChan <- fmt.Errorf("client write error: %v", err)
						}
						return
					}
				}
			}
		}
	}()

	// 等待传输完成或出错
	select {
	case err := <-errChan:
		if err != nil {
			// 只记录非正常关闭的错误
			if !isTimeoutError(err) && !isConnectionClosed(err) {
				pf.logger.Error("Data transfer error (TCP): %v", err)
			}
		}
	case <-ctx.Done():
		pf.logger.Info("Connection closed due to context cancellation")
	}
}

// isTimeoutError 判断是否为超时错误
func isTimeoutError(err error) bool {
	if err == nil {
		return false
	}
	netErr, ok := err.(net.Error)
	return ok && netErr.Timeout()
}

// isConnectionClosed 判断是否为连接关闭错误
func isConnectionClosed(err error) bool {
	if err == nil {
		return false
	}
	return err == io.EOF || strings.Contains(err.Error(), "use of closed network connection") ||
		strings.Contains(err.Error(), "connection reset by peer") ||
		strings.Contains(err.Error(), "broken pipe")
}

// StopForwarding 停止指定端口或所有转发
// 修改 StopForwarding
func (pf *PortForwarder) StopForwarding(listenAddr string) {
	pf.mu.Lock()
	defer pf.mu.Unlock()

	if listenAddr == "all" {
		// 停止所有转发
		for addr := range pf.proxies {
			pf.stopSingleForwarding(addr)
		}
	} else {
		// 停止单个转发
		pf.stopSingleForwarding(listenAddr)
	}
}

// 添加新的辅助方法来处理单个转发的停止
func (pf *PortForwarder) stopSingleForwarding(addr string) {
	if cancel, exists := pf.proxies[addr]; exists {
		// 1. 首先调用 cancel 来通知所有 goroutine 退出
		cancel()

		// 2. 关闭对应的监听器或连接
		if listener, ok := pf.listeners[addr]; ok {
			listener.Close()
			delete(pf.listeners, addr)
		}
		if udpConn, ok := pf.udpConns[addr]; ok {
			udpConn.Close()
			delete(pf.udpConns, addr)
		}

		// 3. 删除代理记录
		delete(pf.proxies, addr)
		pf.logger.Info("Stopped forwarding: %s", addr)
	} else {
		pf.logger.Info("Forwarding not found: %s", addr)
	}
}

// 添加新的方法来检查端口是否已释放
func (pf *PortForwarder) IsPortInUse(addr string) bool {
	// 尝试监听该端口
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return true // 端口仍在使用中
	}
	listener.Close() // 立即关闭测试用的监听器
	return false
}

// Wait 等待所有转发任务完成
func (pf *PortForwarder) Wait() {
	pf.waitGroup.Wait()
	pf.logger.Info("All forwarding tasks completed")
}

// GetActiveForwardings 返回当前活动的转发列表
func (pf *PortForwarder) GetActiveForwardings() []string {
	pf.mu.Lock()
	defer pf.mu.Unlock()

	forwardings := make([]string, 0, len(pf.proxies))
	for addr := range pf.proxies {
		forwardings = append(forwardings, addr)
	}
	return forwardings
}
